# -*- coding: utf-8 -*-
"""
@author: YuHaiyang

"""
import os
from pathlib import Path

import torch
import torchvision.transforms
from PIL import Image

from nets.alex.loader import DataSetLoader

# Press the green button in the gutter to run the script.
if __name__ == '__main__':
    home_path = os.environ.get('HOME')
    src = Path(home_path, "workspace", "dataset", "dc1")

    loader = DataSetLoader(src).gen(1)

    device = torch.device('cuda') if torch.cuda.is_available() else "cpu"

    print("torch.cuda.is_available()", torch.cuda.is_available())

    net = torch.hub.load('pytorch/vision:v0.10.0', 'alexnet', pretrained=True)
    net.eval()

    input_image = Image.open("../../../dataset/dc1/cat/cat.1211.jpg")

    transform = torchvision.transforms.Compose(
        [
            torchvision.transforms.Resize(256),
            torchvision.transforms.CenterCrop(224),
            torchvision.transforms.ToTensor(),
            torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ]
    )

    input_image = transform(input_image)
    print("shape1:", input_image.shape)
    input_image = input_image.unsqueeze(0)
    print("shape2:", input_image.shape)

    with torch.no_grad():
        out_put = net(input_image)

    pred = torch.nn.functional.softmax(out_put[0], dim=0)

    with open("imagenet_classes.txt", "r") as f:
        categories = [s.strip() for s in f.readlines()]

    top5_prob, top5_catid = torch.topk(pred, 5)
    for i in range(top5_prob.size(0)):
        print(categories[top5_catid[i]], top5_prob[i].item())
